"""这段代码实现了一个完整的 RAG（检索增强生成）系统评估流程，主要功能包括文档处理、索引构建、检索增强生成以及系统性能评估,代码的核心功能分为以下几个部分：
1.文档加载与预处理：加载PDF文档并切分为句子或滑动窗口节点。
2.向量索引构建：基于不同节点类型构建两种向量索引（基础句子索引和滑动窗口索引）。
3.检索与问答测试：测试两种索引的检索效果，并输出结果。
4.评估数据集准备：加载预生成的评估问题-答案对
5.多维度评估：使用4种评估器（正确性、相关性、忠实性、语义相似性）对比两种检索器的性能。
"""

import os
import asyncio
import random
import nest_asyncio

# 导入 LlamaIndex 相关模块
from langchain_community.chat_models.zhipuai import ChatZhipuAI
from llama_index.llms.langchain import LangChainLLM
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.core.node_parser import SentenceWindowNodeParser, SentenceSplitter
from llama_index.core import Settings, SimpleDirectoryReader, VectorStoreIndex
from llama_index.core.postprocessor import MetadataReplacementPostProcessor
from llama_index.core.evaluation import (
    BatchEvalRunner, QueryResponseDataset,CorrectnessEvaluator,
    SemanticSimilarityEvaluator, RelevancyEvaluator, FaithfulnessEvaluator
)
from llama_index.core.evaluation.eval_utils import get_responses, get_results_df
from dotenv import load_dotenv
# 加载环境变量
load_dotenv()

# 1. 配置 LLM、Embedding、文本切分器
# --------------------------------------------------
# 设置 LLM（大模型）和 Embedding（向量模型）
"""初始化语言模型（用于生成答案）和嵌入模型（用于文本向量化）。
  关键参数：
    temperature=0.1：控制生成答案的随机性（值越低答案越确定）。
    max_length=512：嵌入模型的最大输入长度。
"""
llm = ChatZhipuAI(api_key=os.getenv("ZHIPUAI_API_KEY"), model="glm-4", temperature=0.1)
embed_model = HuggingFaceEmbedding(
    model_name=r"D:\ideaSpace\MyPython\models\all-MiniLM-L6-v2", max_length=512
)
# SentenceSplitter：用于将文档切分为独立句子节点
text_splitter = SentenceSplitter()
# SentenceWindowNodeParser：用于将文档切分为滑动窗口节点（每个节点包含中心句子及其上下文）
node_parser = SentenceWindowNodeParser.from_defaults(
    window_size=3,
    window_metadata_key="window",
    original_text_metadata_key="original_text",
)
# 全局设置
Settings.llm = llm
Settings.embed_model = embed_model
Settings.text_splitter = text_splitter

# 2. 加载 PDF 文档
# --------------------------------------------------
pdf_path = r"D:\ideaSpace\rag-in-action-master\90-文档-Data\复杂PDF\IPCC_AR6_WGII_Chapter03.pdf"
documents = SimpleDirectoryReader(input_files=[pdf_path]).load_data()

# 3. 文本切分为节点
# --------------------------------------------------
# 滑动窗口节点（用于增强上下文理解）
nodes = node_parser.get_nodes_from_documents(documents)
# 基础句子节点（用于简单检索）
base_nodes = text_splitter.get_nodes_from_documents(documents)

# 4. 构建向量索引
# --------------------------------------------------
"""将切分后的节点转换为可检索的向量索引，区别：
    滑动窗口索引更适合回答需要上下文的复杂问题。
    基础索引适合直接匹配关键词的简单问题
"""
sentence_index = VectorStoreIndex(nodes)
base_index = VectorStoreIndex(base_nodes)

# 5. 检索与问答示例
# --------------------------------------------------
# 滑动窗口检索器
query_engine = sentence_index.as_query_engine(
    # 设置检索返回的顶部结果数量为2
    similarity_top_k=2,
    # 使用元数据替换后处理器：MetadataReplacementPostProcessor用于将检索结果的元数据替换为窗口内容（增强答案连贯性）
    node_postprocessors=[MetadataReplacementPostProcessor(target_metadata_key="window")],
)
# 检索到的答案及其上下文窗口
window_response = query_engine.query("What are the concerns surrounding the AMOC?")
print("\n【滑动窗口检索器返回】\n", window_response)
# 获取第一个检索结果的窗口内容
window = window_response.source_nodes[0].node.metadata["window"]
# 获取第一个检索结果的原始句子文本
sentence = window_response.source_nodes[0].node.metadata["original_text"]
print(f"Window: {window}")
print("------------------")
print(f"Original Sentence: {sentence}")

# 基础检索器
base_query_engine = base_index.as_query_engine(similarity_top_k=2)
vector_response = base_query_engine.query("What are the concerns surrounding the AMOC?")
print("\n【基础检索器返回】\n", vector_response)

# 打印所有 source_nodes 的原始句子
print("\n【滑动窗口检索器 source_nodes 原始句子】")
for source_node in window_response.source_nodes:
    print(source_node.node.metadata["original_text"])
    print("--------")

# 检查基础检索器 source_nodes 是否包含 AMOC
print("\n【基础检索器 source_nodes 是否包含 'AMOC'】")
for node in vector_response.source_nodes:
    print("AMOC mentioned?", "AMOC" in node.node.text)
    print("--------")

# 遍历输出所有基础检索器的 node.text
print("\n【基础检索器 source_nodes 的 node.text】")
for i, node in enumerate(vector_response.source_nodes):
    print(f"第 {i+1} 个 node.text:")
    print(node.node.text)
    print("--------")

# 6. 评测数据集准备
# --------------------------------------------------
# 采样部分节点用于生成评测问题
num_nodes_eval = 30
sample_eval_nodes = random.sample(base_nodes[:200], num_nodes_eval)
# 生成评测数据集（如未生成可取消注释）
# dataset_generator = DatasetGenerator(
#     sample_eval_nodes,
#     llm=OpenAI(model="gpt-4"),
#     show_progress=True,
#     num_questions_per_chunk=2,
# )
# eval_dataset = await dataset_generator.agenerate_dataset_from_nodes()
# eval_dataset.save_json("90-文档-Data/复杂PDF/ipcc_eval_qr_dataset.json")

# 加载已生成的评测数据集
eval_dataset = QueryResponseDataset.from_json(r"D:\ideaSpace\rag-in-action-master\90-文档-Data\复杂PDF\ipcc_eval_qr_dataset.json")

# 7. 构建评测器
# --------------------------------------------------
"""
正确性：答案是否与参考答案一致。
相关性：答案是否与问题相关。
忠实性：答案是否严格基于检索内容（避免幻觉）。
语义相似性：答案与参考答案的语义匹配度（基于嵌入模型）。
"""
evaluator_c = CorrectnessEvaluator(llm=LangChainLLM(ChatZhipuAI(api_key=os.getenv("ZHIPUAI_API_KEY"))))  # 正确性
evaluator_s = SemanticSimilarityEvaluator() # 语义相似性
evaluator_r = RelevancyEvaluator(llm=LangChainLLM(ChatZhipuAI(api_key=os.getenv("ZHIPUAI_API_KEY")))) # 相关性
evaluator_f = FaithfulnessEvaluator(llm=LangChainLLM(ChatZhipuAI(api_key=os.getenv("ZHIPUAI_API_KEY"), model="glm-4"))) # 忠实性
# pairwise_evaluator = PairwiseComparisonEvaluator(llm=OpenAI(model="gpt-4"))

evaluator_dict = {
    "correctness": evaluator_c,
    "faithfulness": evaluator_f,
    "relevancy": evaluator_r,
    "semantic_similarity": evaluator_s,
}
batch_runner = BatchEvalRunner(evaluator_dict, workers=2, show_progress=True)

# 8. 评测主流程
# --------------------------------------------------
async def main():
    max_samples = 30
    eval_qs = eval_dataset.questions
    ref_response_strs = [r for (_, r) in eval_dataset.qr_pairs]

    # 重新构建检索器，保证评测一致
    base_query_engine = base_index.as_query_engine(similarity_top_k=2)
    window_query_engine = sentence_index.as_query_engine(
        similarity_top_k=2,
        node_postprocessors=[MetadataReplacementPostProcessor(target_metadata_key="window")],
    )

    # 对两种检索器的答案进行多维度评估，生成对比结果表
    # 获取模型回答
    base_pred_responses = get_responses(
        eval_qs[:max_samples], base_query_engine, show_progress=True
    )
    pred_responses = get_responses(
        eval_qs[:max_samples], window_query_engine, show_progress=True
    )

    # 评测
    eval_results = await batch_runner.aevaluate_responses(
        queries=eval_qs[:max_samples],
        responses=pred_responses[:max_samples],
        reference=ref_response_strs[:max_samples],
    )
    base_eval_results = await batch_runner.aevaluate_responses(
        queries=eval_qs[:max_samples],
        responses=base_pred_responses[:max_samples],
        reference=ref_response_strs[:max_samples],
    )
    results_df = get_results_df(
        [eval_results, base_eval_results],
        ["Sentence Window Retriever", "Base Retriever"],
        ["correctness", "relevancy", "faithfulness", "semantic_similarity"],
    )
    print("\n【评测结果】")
    print(results_df)

if __name__ == "__main__":
    nest_asyncio.apply()
    asyncio.run(main())

